import logging
from typing import override

logger = logging.getLogger(__name__)

from app.retriever.base_retriever import Retriever, CustomTextSplitter, docs_resolve


class FaqRetriever(Retriever):
    def __init__(self, collection_name="knowledge"):
        super().__init__(collection_name)
        # 这里只需要使用简单的字符分割器即可，因为数据相对规整
        self.text_splitter = CustomTextSplitter(chunk_overlap=0, chunk_size=100)
        self.docs_resolver = docs_resolve

    # 加载形如question|answer的文本
    @override
    def load_data_from_txt(self, data_dir: str = "../data", suffix: str = ".txt"):
        super().load_data_from_txt(data_dir, suffix)

    # 加载标准的json数组
    @override
    def load_data_from_json(self, data_dir: str = "../data", suffix: str = ".json"):
        super().load_data_from_json(data_dir, suffix)


if __name__ == '__main__':
    retriever = FaqRetriever()
    retriever.load_data_from_json(data_dir="../data")
